/*
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You under the Apache License, Version 2.0
 * (the "License"); you may not use this file except in compliance with
 * the License.  You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package com.lvyh.lightframe.sharding.client.interceptor;

import com.lvyh.lightframe.sharding.client.common.DataTableInfo;
import com.lvyh.lightframe.sharding.client.common.ShardModel;
import com.lvyh.lightframe.sharding.client.context.MultiDataSourceHolder;
import com.lvyh.lightframe.sharding.client.context.ShardingContext;
import com.lvyh.lightframe.sharding.client.datasource.DynamicRoutingDataSource;
import com.lvyh.lightframe.sharding.client.enums.ShardingStrategyType;
import com.lvyh.lightframe.sharding.client.parser.SimpleSqlParser;
import com.lvyh.lightframe.sharding.client.strategy.ShardingStrategyFactory;
import com.lvyh.lightframe.sharding.client.util.CommonUtils;
import com.lvyh.lightframe.sharding.client.util.ComparatorUtil;
import com.lvyh.lightframe.sharding.client.util.ExecutorUtil;
import com.lvyh.lightframe.sharding.client.util.Page;
import com.lvyh.lightframe.sharding.client.util.PageContext;
import com.lvyh.lightframe.sharding.client.util.ReflectionUtil;
import com.lvyh.lightframe.sharding.client.util.SpringContextUtils;
import com.lvyh.lightframe.sharding.client.util.SqlUtils;
import org.apache.ibatis.cache.CacheKey;
import org.apache.ibatis.executor.Executor;
import org.apache.ibatis.mapping.BoundSql;
import org.apache.ibatis.mapping.MappedStatement;
import org.apache.ibatis.mapping.SqlCommandType;
import org.apache.ibatis.plugin.Interceptor;
import org.apache.ibatis.plugin.Intercepts;
import org.apache.ibatis.plugin.Invocation;
import org.apache.ibatis.plugin.Plugin;
import org.apache.ibatis.plugin.Signature;
import org.apache.ibatis.reflection.DefaultReflectorFactory;
import org.apache.ibatis.reflection.MetaObject;
import org.apache.ibatis.reflection.ReflectorFactory;
import org.apache.ibatis.reflection.factory.DefaultObjectFactory;
import org.apache.ibatis.reflection.factory.ObjectFactory;
import org.apache.ibatis.reflection.wrapper.DefaultObjectWrapperFactory;
import org.apache.ibatis.reflection.wrapper.ObjectWrapperFactory;
import org.apache.ibatis.session.ResultHandler;
import org.apache.ibatis.session.RowBounds;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;
import org.springframework.util.CollectionUtils;
import org.springframework.util.StringUtils;

import java.util.ArrayList;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.Properties;

/**
 * @author lvyh 2021/10/28.
 */
@Component
@Intercepts({
        @Signature(type = Executor.class, method = "query", args = {MappedStatement.class, Object.class, RowBounds.class, ResultHandler.class}),
        @Signature(type = Executor.class, method = "query", args = {MappedStatement.class, Object.class, RowBounds.class, ResultHandler.class, CacheKey.class, BoundSql.class})
})
public class SelectInterceptor implements Interceptor {
    private Logger logger = LoggerFactory.getLogger(this.getClass());
    private static final ObjectFactory DEFAULT_OBJECT_FACTORY = new DefaultObjectFactory();
    private static final ObjectWrapperFactory DEFAULT_OBJECT_WRAPPER_FACTORY = new DefaultObjectWrapperFactory();
    private static final ReflectorFactory REFLECTOR_FACTORY = new DefaultReflectorFactory();

    @Autowired
    private ShardingStrategyFactory shardingStrategyFactory;

    private Map<String, MappedStatement> mappedStatementMap = new HashMap<>();

    @Override
    public Object intercept(Invocation invocation) throws Throwable {
        Object[] args = invocation.getArgs();
        MappedStatement ms = (MappedStatement) args[0];
        Object parameter = args[1];
        RowBounds rowBounds = (RowBounds) args[2];
        ResultHandler resultHandler = (ResultHandler) args[3];
        Executor executor = (Executor) invocation.getTarget();
        CacheKey cacheKey;
        BoundSql boundSql;

        String sqlCommandType = ms.getSqlCommandType().name();
        if (!Objects.equals(SqlCommandType.SELECT.name(), sqlCommandType)) {
            return invocation.proceed();
        }

        if (args.length == 4) {
            boundSql = ms.getBoundSql(parameter);
            cacheKey = executor.createCacheKey(ms, parameter, rowBounds, boundSql);
            logger.info("[SelectInterceptor] print msId:{}", ms.getId());
        } else {
            cacheKey = (CacheKey) args[4];
            boundSql = (BoundSql) args[5];
        }
        try {
            ShardModel shardModel = ShardingContext.get();
            if (shardModel != null) {
                return doShardQuery(executor, rowBounds, resultHandler, ms, boundSql, cacheKey, shardModel, parameter);
            } else {
                Page page = PageContext.getPage();
                if (page != null) {
                    return this.simplePageQuery(page, boundSql.getSql(), executor, ms, boundSql, cacheKey, rowBounds, resultHandler, parameter);
                }
                return invocation.proceed();
            }
        } finally {
            PageContext.clear();
        }
    }

    private List<Object> simplePageQuery(Page page, String sql, Executor executor, MappedStatement ms, BoundSql boundSql, CacheKey cacheKey, RowBounds rowBounds, ResultHandler resultHandler, Object parameter) throws Exception {
        List<Object> list = this.queryByPage(page, sql, executor, ms, boundSql, cacheKey, rowBounds, resultHandler, parameter);
        return list;
    }

    private List<Object> doShardQuery(Executor executor, RowBounds rowBounds, ResultHandler resultHandler, MappedStatement ms, BoundSql boundSql, CacheKey cacheKey, ShardModel shardModel, Object parameter) throws Exception {
        List<Object> list = shardQuery(executor, rowBounds, resultHandler, ms, boundSql, cacheKey, shardModel, parameter);
        return list;
    }


    private List<Object> shardQuery(Executor executor, RowBounds rowBounds, ResultHandler resultHandler, MappedStatement ms, BoundSql boundSql, CacheKey cacheKey, ShardModel shardModel, Object parameter) throws Exception {
        String shardingColumn = shardModel.getShardKey();

        String sqlId = ms.getId();
        String shardingValue = null;
        if (sqlId.endsWith("selectByExample_COUNT") || sqlId.endsWith("selectByExample")) {
            shardingValue = computeShardingValue(shardingColumn, parameter);
        } else
            shardingValue = (String) ReflectionUtil.getFieldValue(parameter, shardingColumn);

        if (!StringUtils.isEmpty(shardingValue)) {
            //shard key is not empty
            return queryShardKeyValueExist(shardingValue, executor, rowBounds, resultHandler, ms, boundSql, cacheKey, shardModel, parameter);
        } else {
            //shard key is empty and alg is Hash
            return queryShardKeyValueNotExist(executor, rowBounds, resultHandler, ms, boundSql, cacheKey, shardModel, parameter);
        }
    }

    private List<Object> queryShardKeyValueNotExist(Executor executor, RowBounds rowBounds, ResultHandler resultHandler, MappedStatement ms, BoundSql boundSql, CacheKey cacheKey, ShardModel shardModel, Object parameter) throws Exception {
        Page page = PageContext.getPage();
        String shardingAlg = shardModel.getShardAlg();
        int dbShardingNum = shardModel.getDbShardNum();
        int tableShardingNum = shardModel.getTableShardNum();
        String logicTableName = shardModel.getLogicTableName();
        String sql = boundSql.getSql();
        List<Object> result = new ArrayList<>();
        if (!Objects.equals(shardingAlg, ShardingStrategyType.HASH.getValue())) {
            throw new Exception("sharding strategy type is invalid!");
        }

        if (page != null) {
            long total = 0;
            for (int i = 0; i < tableShardingNum; i++) {
                String newSql = null;
                if (dbShardingNum > 1 && tableShardingNum == 1) {
                    newSql = sql;
                } else {
                    newSql = SimpleSqlParser.getConvertedSql(logicTableName, i, sql);
                    ReflectionUtil.setField(boundSql, "sql", newSql);
                }

                Long count = this.getCount(newSql, executor, ms, parameter, rowBounds, resultHandler, boundSql);
                total += count;
                if (count > 0) {
                    cacheKey.update(new Object());
                    String limitSql = getNewLimitSql(page, newSql, true);
                    List<Object> query = ExecutorUtil.query(limitSql, executor, ms, rowBounds, resultHandler, boundSql, cacheKey);
                    result.addAll(query);
                }
            }

            page.setTotal(total);
            if (page.getTotal() > 0) {
                if (page.getOrderBy() != null) {
                    String[] arr = page.getOrderBy().split(" ");//update_time desc
                    String orderColumn = CommonUtils.convertColumn(arr[0]);
                    boolean isAsc = arr[1].equals("asc") ? true : false;
                    Collections.sort(result, new ComparatorUtil(orderColumn, isAsc));
                }
                return page.page(result);
            }
        } else {
            for (int i = 0; i < tableShardingNum; i++) {
                String newSql = null;
                if (dbShardingNum > 1 && tableShardingNum == 1) {
                    newSql = sql;
                } else {
                    newSql = SimpleSqlParser.getConvertedSql(logicTableName, i, sql);
                    ReflectionUtil.setField(boundSql, "sql", newSql);
                }

                cacheKey.update(new Object());
                List<Object> queryResult = ExecutorUtil.query(newSql, executor, ms, rowBounds, resultHandler, boundSql, cacheKey);
                result.addAll(queryResult);
            }
        }

        return result;
    }

    private List<Object> queryShardKeyValueExist(String shardingValue, Executor executor, RowBounds rowBounds, ResultHandler resultHandler, MappedStatement ms, BoundSql boundSql, CacheKey cacheKey, ShardModel shardModel, Object parameter) throws Exception {
        String logicDsKey = shardModel.getLogicDatasourceName();
        String logicTableName = shardModel.getLogicTableName();
        String shardingAlg = shardModel.getShardAlg();
        int tableShardingNum = shardModel.getTableShardNum();
        int dbShardingNum = shardModel.getDbShardNum();

        DataTableInfo dataTableInfo = this.shardingStrategyFactory.getShardInfo(shardingAlg, logicDsKey, logicTableName, shardingValue, tableShardingNum, dbShardingNum);
        String sql = boundSql.getSql();
        Page page = PageContext.getPage();
        List<Object> result;
        String newSql = SimpleSqlParser.getConvertedSql(logicTableName, dataTableInfo.getPhysicalTableName(), sql);
        ReflectionUtil.setField(boundSql, "sql", newSql);

        if (page != null) {
            result = queryByPage(page, newSql, executor, ms, boundSql, cacheKey, rowBounds, resultHandler, parameter);
        } else {
            result = ExecutorUtil.query(newSql, executor, ms, rowBounds, resultHandler, boundSql, cacheKey);
        }
        return result;
    }


    private String getNewLimitSql(Page page, String sql, boolean crossDbQuery) {
        StringBuffer sb = new StringBuffer();
        sb.append(sql);
        //update_time desc
        if (page.getOrderBy() != null) {
            String[] arr = page.getOrderBy().split(" ");
            String orderColumn = arr[0];
            String orderType = arr[1];
            sb.append(" ORDER BY ");
            sb.append(" " + orderColumn);
            sb.append(" " + orderType);
        }

        sb.append(" limit ");

        //page num>1 ,will rewrite sql: limit 0 , page_num * page_size
        if (crossDbQuery) {
            sb.append("0");
            sb.append("," + page.getPageNum() * page.getPageSize());
        } else {
            sb.append(page.getStart());
            sb.append("," + page.getPageSize());
        }

        return sb.toString();
    }

    private List<Object> queryByPage(Page page, String sql, Executor executor, MappedStatement ms, BoundSql boundSql, CacheKey cacheKey, RowBounds rowBounds, ResultHandler resultHandler, Object parameter) throws Exception {
        Long count = getCount(sql, executor, ms, parameter, rowBounds, resultHandler, boundSql);
        page.setTotal(count.intValue());

        String newLimitSql = getNewLimitSql(page, sql, false);
        List<Object> result = ExecutorUtil.query(newLimitSql, executor, ms, rowBounds, resultHandler, boundSql, cacheKey);
        page.addAll(CollectionUtils.isEmpty(result) ? new ArrayList() : result);
        return page;
    }

    private Long getCount(String sql, Executor executor, MappedStatement mappedStatement, Object parameter, RowBounds rowBounds, ResultHandler resultHandler, BoundSql boundSql) throws Exception {
        //Build a new count query statement
        String countMsId = mappedStatement.getId() + "_COUNT";
        MappedStatement countStatement = ExecutorUtil.getExistedMappedStatement(mappedStatement.getConfiguration(), countMsId);
        Long count;
        if (countStatement != null) {
            count = ExecutorUtil.executeManualCount(executor, countStatement, parameter, boundSql, resultHandler);
        } else {
            countStatement = mappedStatementMap.get(countMsId);
            if (countStatement == null) {
                countStatement = SqlUtils.newCountMappedStatement(mappedStatement, countMsId);
                mappedStatementMap.put(countMsId, countStatement);
            }
            count = ExecutorUtil.executeAutoCount(sql, executor, countStatement, parameter, boundSql, rowBounds, resultHandler);
        }
        return count;
    }

    private MetaObject getMetaObject(Object parameterObject) {
        MetaObject metaObject = MetaObject.forObject(parameterObject, DEFAULT_OBJECT_FACTORY, DEFAULT_OBJECT_WRAPPER_FACTORY, REFLECTOR_FACTORY);
        return metaObject;
    }

    private String computeShardingValue(String shardingColumn, Object parameterObject) {
        // Reflection get parameter object
        MetaObject param = getMetaObject(parameterObject);
        Object oredCriteriaList = param.getValue("oredCriteria");
        if (oredCriteriaList != null && oredCriteriaList instanceof ArrayList) {
            Object criteria = ((ArrayList) oredCriteriaList).get(0);
            MetaObject criteriaBaseObject = getMetaObject(criteria);
            Object criteriaList = criteriaBaseObject.getValue("criteria");
            if (criteriaList != null && criteriaList instanceof ArrayList) {
                if (!CollectionUtils.isEmpty(((ArrayList) criteriaList))) {
                    Object criteriaObject = ((ArrayList) criteriaList).get(0);
                    MetaObject criteriaMetaObject = getMetaObject(criteriaObject);
                    Object condition = criteriaMetaObject.getValue("condition");
                    Object value = criteriaMetaObject.getValue("value");
                    if (condition instanceof String && condition.toString().contains(shardingColumn)) {
                        return value.toString();
                    }
                } else {
                    return "";//shard key value is null
                }
            }
        }
        return "";
    }

    @Override
    public Object plugin(Object target) {
        return Plugin.wrap(target, this);
    }

    @Override
    public void setProperties(Properties properties) {

    }
}
